#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy, sys
import moveit_commander

import tf.transformations as tfs
import numpy as np
import math

import kinematics as ks
from arm_manager import ArmManager
from gripper_manager import GripperManager
#-------------------------------------
# from chassis import ChassisManager
#-------------------------------------
from sensors_init import SensorsInit

from moveit_msgs.msg import RobotTrajectory
from trajectory_msgs.msg import JointTrajectoryPoint
from geometry_msgs.msg import PoseStamped, Pose
from std_msgs.msg import UInt8
from std_msgs.msg import Float64MultiArray,Float32MultiArray
from webots_ros.msg import RecognitionObject


class GraspDemo:
    def __init__(self):
        # 初始化move_group的API
        moveit_commander.roscpp_initialize(sys.argv)
        # 初始化ROS节点
        rospy.init_node('moveit_fk_demo')
       
        self.setonceflag = 1
        self.enable_flag = 0
        self.objects_info = ['none','can','model']  
        self.camera_to_base_R_T = None
        self.precal_camera_to_base_R_T()
        
        #arm
        self.cmarm = ArmManager()
        #gripper
        self.cmgripper = GripperManager()
#-------------------------------------------------------
        #chassis
        # self.cm = ChassisManager()
        self.cmd_publisher = rospy.Publisher('/robot/cmd', Float32MultiArray, queue_size=1)
        # cmd_info = Float32MultiArray()
        # cmd_publisher.publish(std_msgs.msg.String("foo"))
#-------------------------------------------------------
        #camera
        self.si = SensorsInit()
        self.si.enable_sensors()

        rospy.Subscriber('/sign', UInt8, self.get_sign_callback)
        rospy.Subscriber('/robot/kinect_color/recognition_objects', RecognitionObject, self.get_object_info_callback)
        rospy.spin()

    def precal_camera_to_base_R_T(self):
            camera_to_chassis_base_R = tfs.euler_matrix(-0.261795,3.1415926,0,'sxyz')
            camera_to_chassis_base_T = (np.array([[0,1.08+0.05,0,1]])).T
            for i in range(3):
                camera_to_chassis_base_R[i,3] = camera_to_chassis_base_T[i,0]

            chassis_base_to_arm_R = (tfs.euler_matrix(1.5708, 0, 1.5708)).T
            chassis_base_to_arm_T = -np.dot(chassis_base_to_arm_R,(np.array([[0,0.06,0,1]])).T)
            for i in range(3):
                chassis_base_to_arm_R[i,3] = chassis_base_to_arm_T[i,0]

            self.camera_to_base_R_T = np.dot(chassis_base_to_arm_R,camera_to_chassis_base_R)

    def get_sign_callback(self,msg_sign):
        self.enable_flag = msg_sign.data
        cmd_info = Float32MultiArray()
        cmd_info.data = np.array([1,0,4])
        for i in range(5):
            rospy.sleep(0.05)
            self.cmd_publisher.publish(cmd_info)

    def get_object_info_callback(self,msg_object_info):
        if(self.objects_info[self.enable_flag] == msg_object_info.model and self.enable_flag):
        # if(self.objects_info[1] == msg_object_info.model):
            # 流程1： 横向移动，动态调整
            if(abs(msg_object_info.position.x+0.015)>0.01):
                if(self.setonceflag):
                #----------调整朝向------------
                    cmd_info = Float32MultiArray()
                    cmd_info.data = [1,0.0,-np.sign(msg_object_info.position.x)*4] 
                    print(-np.sign(msg_object_info.position.x)*4,msg_object_info.position.x)
                    for i in range(5):
                        rospy.sleep(0.05)
                        self.cmd_publisher.publish(cmd_info)
                    self.setonceflag = 0
            #流程2： 停止后，执行静态抓取，不带调整
            else:
                #-----------停止调整----------
                cmd_info = Float32MultiArray()
                cmd_info.data = np.array([1,0,0])
                for i in range(5):
                    rospy.sleep(0.05)
                    self.cmd_publisher.publish(cmd_info)
                rospy.sleep(2)

                ob_to_camera = (np.array([[msg_object_info.position.x,msg_object_info.position.y,msg_object_info.position.z,1]])).T
                ob_to_base = np.dot(self.camera_to_base_R_T,ob_to_camera)
                print("--------------------------------")
                print(ob_to_base)
                self.set_pose(ob_to_base)
                self.enable_flag = 0

    def set_pose(self,taget_pose):
#---------------抬起机械臂--------------
        joint_positions,x_offset = ks.IK_with_chassis(taget_pose[1,0],taget_pose[0,0],0)
        print(joint_positions)
        self.cmarm.set_velocity([joint_positions[0],joint_positions[1],0,joint_positions[2],0],1)
        rospy.sleep(1)
#---------------轮式里程计向前走--------------
        cmd_info = Float32MultiArray()
        cmd_info.data = [2,x_offset,1] 
        for i in range(5):
            rospy.sleep(0.05)
            self.cmd_publisher.publish(cmd_info)
        rospy.sleep(x_offset/1)
#---------------夹取--------------
        self.cmgripper.set_velocity(0.8/0.05*(0.05-0.027),0.15)
        rospy.sleep(8)
# #--------------抬起动作-----------------
        joint_positions,x_offset = ks.IK_with_chassis(taget_pose[1,0],taget_pose[0,0]+0.1,0)
        self.cmarm.set_velocity([joint_positions[0],joint_positions[1],0,joint_positions[2],0],0.2)
        rospy.sleep(2)

if __name__ == "__main__":
    try:
        GraspDemo()
    except rospy.ROSInterruptException:
        pass

